package com.lnn.springBase.utils.image;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.util.ArrayList;
import java.util.List;

public class PrecipitationVisualizer {

    // 中国地理范围常量
    private static final double CHINA_MIN_LON = 73.66;
    private static final double CHINA_MAX_LON = 135.05;
    private static final double CHINA_MIN_LAT = 18.15;
    private static final double CHINA_MAX_LAT = 53.55;

    // 默认颜色分段（降水强度）
    private static final Color[] PRECIP_COLORS = {
            new Color(240, 249, 255),  // 0-10mm
            new Color(189, 215, 238),  // 10-25mm
            new Color(107, 174, 214),  // 25-50mm
            new Color(49, 130, 189),   // 50-100mm
            new Color(8, 81, 156)      // >100mm
    };

    public static class VisualConfig {
        public int imageWidth = 800;
        public int imageHeight = 600;
        public String title = "全国降水量分布图";
        public Font titleFont = new Font("宋体", Font.BOLD, 24);
        public Color titleColor = Color.BLACK;
        public int legendWidth = 300;
        public int legendHeight = 30;
        public boolean drawBorder = true;
    }

    /**
     * 生成降水可视化图
     * @param lons 经度数组（从西到东）
     * @param lats 纬度数组（从南到北）
     * @param precip 降水数据[lon][lat]
     * @param config 可视化配置
     */
    public static BufferedImage createPrecipitationMap(
            double[] lons,
            double[] lats,
            double[][] precip,
            VisualConfig config) {

        BufferedImage image = new BufferedImage(
                config.imageWidth,
                config.imageHeight,
                BufferedImage.TYPE_INT_ARGB);

        Graphics2D g = image.createGraphics();
        g.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
                RenderingHints.VALUE_ANTIALIAS_ON);

        // 绘制背景
        g.setColor(Color.WHITE);
        g.fillRect(0, 0, config.imageWidth, config.imageHeight);

        // 绘制降水数据
        drawPrecipitation(g, lons, lats, precip, config);

        // 绘制标题
        drawTitle(g, config);

        // 绘制图例
        drawLegend(g, config);

        // 绘制国界
        if (config.drawBorder) {
            drawChinaBorder(g, config);
        }

        g.dispose();
        return image;
    }

    private static void drawPrecipitation(Graphics2D g,
                                          double[] lons,
                                          double[] lats,
                                          double[][] precip,
                                          VisualConfig config) {

        int lonCount = lons.length;
        int latCount = lats.length;

        for (int i = 0; i < lonCount; i++) {
            for (int j = 0; j < latCount; j++) {
                double value = precip[i][j];
                if (Double.isNaN(value)) continue;

                // 转换为图像坐标
                Point p = geoToPixel(lons[i], lats[j], config);

                // 获取颜色
                Color color = getPrecipColor(value);

                // 绘制像素
                g.setColor(color);
                g.fillRect(p.x, p.y, 1, 1);
            }
        }
    }

    private static Point geoToPixel(double lon, double lat, VisualConfig config) {
        double xRatio = (lon - CHINA_MIN_LON) / (CHINA_MAX_LON - CHINA_MIN_LON);
        double yRatio = 1 - (lat - CHINA_MIN_LAT) / (CHINA_MAX_LAT - CHINA_MIN_LAT);

        return new Point(
                (int) (xRatio * config.imageWidth),
                (int) (yRatio * config.imageHeight)
        );
    }

    private static Color getPrecipColor(double value) {
        if (value < 10) return PRECIP_COLORS[0];
        if (value < 25) return PRECIP_COLORS[1];
        if (value < 50) return PRECIP_COLORS[2];
        if (value < 100) return PRECIP_COLORS[3];
        return PRECIP_COLORS[4];
    }

    private static void drawTitle(Graphics2D g, VisualConfig config) {
        g.setFont(config.titleFont);
        g.setColor(config.titleColor);

        FontMetrics fm = g.getFontMetrics();
        int titleWidth = fm.stringWidth(config.title);

        int x = (config.imageWidth - titleWidth) / 2;
        int y = fm.getAscent() + 10;

        g.drawString(config.title, x, y);
    }

    private static void drawLegend(Graphics2D g, VisualConfig config) {
        int legendX = config.imageWidth - config.legendWidth - 20;
        int legendY = config.imageHeight - 50;

        // 绘制颜色条
        for (int i = 0; i < PRECIP_COLORS.length; i++) {
            int x = legendX + i * (config.legendWidth / PRECIP_COLORS.length);
            int width = config.legendWidth / PRECIP_COLORS.length;

            g.setColor(PRECIP_COLORS[i]);
            g.fillRect(x, legendY, width, 20);
        }

        // 绘制文字标注
        g.setColor(Color.BLACK);
        g.setFont(new Font("宋体", Font.PLAIN, 12));
        String[] labels = {"0-10", "10-25", "25-50", "50-100", ">100"};
        for (int i = 0; i <= labels.length; i++) {
            int x = legendX + i * (config.legendWidth / labels.length);
            String label = (i < labels.length) ? labels[i] : "(mm)";
            g.drawString(label, x - 10, legendY + 35);
        }
    }

    private static void drawChinaBorder(Graphics2D g, VisualConfig config) {
        // 此处应添加实际地理边界数据，以下是简化示例
        g.setColor(Color.BLACK);
        g.setStroke(new BasicStroke(1));

        // 示例绘制矩形边界
        Point sw = geoToPixel(CHINA_MIN_LON, CHINA_MIN_LAT, config);
        Point ne = geoToPixel(CHINA_MAX_LON, CHINA_MAX_LAT, config);
        g.drawRect(sw.x, sw.y, ne.x - sw.x, ne.y - sw.y);
    }

    public static void main(String[] args) throws Exception {
        // 生成示例数据
        double[] lons = createRange(CHINA_MIN_LON, CHINA_MAX_LON, 0.5);
        double[] lats = createRange(CHINA_MIN_LAT, CHINA_MAX_LAT, 0.5);
        double[][] precip = generateSampleData(lons.length, lats.length);

        VisualConfig config = new VisualConfig();
        config.title = "2023年7月全国降水量分布";

        BufferedImage image = createPrecipitationMap(lons, lats, precip, config);

        // 保存图片
        javax.imageio.ImageIO.write(image, "PNG", new java.io.File("precipitation.png"));
    }

    private static double[] createRange(double start, double end, double step) {
        List<Double> list = new ArrayList<>();
        for (double v = start; v <= end; v += step) {
            list.add(v);
        }
        return list.stream().mapToDouble(Double::doubleValue).toArray();
    }

    private static double[][] generateSampleData(int width, int height) {
        double[][] data = new double[width][height];
        for (int i = 0; i < width; i++) {
            for (int j = 0; j < height; j++) {
                data[i][j] = Math.random() * 150; // 0-150mm随机数据
                if (Math.random() < 0.2) data[i][j] = Double.NaN; // 20%缺测
            }
        }
        return data;
    }
}
